package com.kasax.plugin.shadow.transform

import com.android.build.api.transform.DirectoryInput
import com.android.build.api.transform.Format
import com.android.build.api.transform.JarInput
import com.android.build.api.transform.QualifiedContent
import com.android.build.api.transform.Transform
import com.android.build.api.transform.TransformException
import com.android.build.api.transform.TransformInput
import com.android.build.api.transform.TransformInvocation
import com.android.build.api.transform.TransformOutputProvider
import com.android.build.gradle.internal.pipeline.TransformManager
import com.android.utils.FileUtils
import com.kasax.plugin.shadow.ShadowExtenstion
import com.kasax.plugin.shadow.utils.LogUtils
import com.kasax.plugin.shadow.visitor.LifecycleClassVisitor
import com.kasax.plugin.shadow.visitor.PackageRemappingClassVisitor
import org.apache.commons.io.IOUtils
import org.objectweb.asm.ClassReader
import org.objectweb.asm.ClassVisitor
import org.objectweb.asm.ClassWriter

import java.util.jar.JarEntry
import java.util.jar.JarFile
import java.util.jar.JarOutputStream
import java.util.zip.ZipEntry

class TestTransform extends Transform {

    private static final String UI_PACKAGE = "com/kasax/testsdk/ui/"
    private ShadowExtenstion shadowExtenstion

    public TestTransform(ShadowExtenstion shadowExtenstion) {
        this.shadowExtenstion = shadowExtenstion
    }

    @Override
    String getName() {
        return "TestTransform"
    }

    @Override
    Set<QualifiedContent.ContentType> getInputTypes() {
        return TransformManager.CONTENT_CLASS
    }

    @Override
    Set<? super QualifiedContent.Scope> getScopes() {
        return TransformManager.PROJECT_ONLY
    }

    @Override
    boolean isIncremental() {
        return true
    }

    @Override
    void transform(TransformInvocation transformInvocation) throws TransformException, InterruptedException, IOException {
        super.transform(transformInvocation)
        LogUtils.d(getName(), "------------ transform start ------------")
        showRelocPkgs()

        long startTime = System.currentTimeMillis()
        TransformOutputProvider outputProvider = transformInvocation.outputProvider
        if (outputProvider != null) {
            outputProvider.deleteAll()
        }

        transformInvocation.inputs.each { TransformInput input ->
            input.jarInputs.each { JarInput jarInput ->
                handleJarInput(jarInput, outputProvider)
            }
            input.directoryInputs.each { DirectoryInput directoryInput ->
                handleDirInput(directoryInput, outputProvider)
            }
        }

        long cost = System.currentTimeMillis() - startTime
        LogUtils.d(getName(), "------------ transform complete, cost: ${cost}ms ------------")
    }

    private void showRelocPkgs() {
        if (shadowExtenstion != null && shadowExtenstion.relocPkgs != null && !shadowExtenstion.relocPkgs.isEmpty()) {
            shadowExtenstion.relocPkgs.entrySet().each {
                LogUtils.d(getName(), "o-pkg: ${it.key} -> new-pkg: ${it.value}")
            }
        }
    }

    private void handleJarInput(JarInput jarInput, TransformOutputProvider outputProvider) {
        File tempFile = new File(jarInput.file.parent, "${jarInput.file.name}.temp")
        if (tempFile.exists()) {
            tempFile.delete()
        }

        JarFile jarFile = new JarFile(jarInput.file)
        JarOutputStream jarOs = new JarOutputStream(new FileOutputStream(tempFile))

        jarFile.entries().each { JarEntry jarEntry ->
            String entryName = jarEntry.name
            InputStream ins = jarFile.getInputStream(jarEntry)
            jarOs.putNextEntry(new ZipEntry(entryName))

            if (entryName.endsWith(".class") && !entryName.startsWith("R\$") && !entryName.equals("R.class") && !entryName.equals("BuildConfig.class")) {
                byte[] sourceBytes = IOUtils.toByteArray(ins)
                ClassReader classReader = new ClassReader(sourceBytes)
                ClassWriter classWriter = new ClassWriter(classReader, ClassWriter.COMPUTE_MAXS)

                ClassVisitor visitor = new PackageRemappingClassVisitor(classWriter,shadowExtenstion)
//                if (entryName.startsWith(UI_PACKAGE)) {
//                    LogUtils.d(getName(), "Applying LifecycleClassVisitor for $entryName")
//                    visitor = new LifecycleClassVisitor(visitor)
//                }

                classReader.accept(visitor, ClassReader.EXPAND_FRAMES)
                jarOs.write(classWriter.toByteArray())
            } else {
                jarOs.write(IOUtils.toByteArray(ins))
            }
            jarOs.closeEntry()
        }

        jarOs.close()
        jarFile.close()

        File dest = outputProvider.getContentLocation(jarInput.name, jarInput.contentTypes, jarInput.scopes, Format.JAR)
        FileUtils.copyFile(tempFile, dest)
        tempFile.delete()
    }

    private void handleDirInput(DirectoryInput directoryInput, TransformOutputProvider outputProvider) {
        File dest = outputProvider.getContentLocation(directoryInput.name, directoryInput.contentTypes, directoryInput.scopes, Format.DIRECTORY)
        FileUtils.copyDirectory(directoryInput.file, dest)

        dest.eachFileRecurse { File file ->
            String fileName = file.name
            if (fileName.endsWith(".class") && !fileName.startsWith("R\$") && !fileName.equals("R.class") && !fileName.equals("BuildConfig.class")) {
                try {
                    byte[] sourceBytes = file.bytes
                    ClassReader classReader = new ClassReader(sourceBytes)
                    ClassWriter classWriter = new ClassWriter(classReader, ClassWriter.COMPUTE_MAXS)

                    ClassVisitor visitor = new PackageRemappingClassVisitor(classWriter, shadowExtenstion)
//                    String className = file.absolutePath.replace(dest.absolutePath + File.separator, "").replace(File.separator, "/")
//                    if (className.startsWith(UI_PACKAGE)) {
//                        LogUtils.d(getName(), "Applying LifecycleClassVisitor for $className")
//                        visitor = new LifecycleClassVisitor(visitor)
//                    }

                    classReader.accept(visitor, ClassReader.EXPAND_FRAMES)
                    byte[] modifiedBytes = classWriter.toByteArray()
                    if (modifiedBytes.length > 0) {
                        file.withOutputStream { it.write(modifiedBytes) }
                    } else {
                        LogUtils.e(getName(), "Generated bytecode is empty for file: $fileName")
                    }
                } catch (Exception e) {
                    LogUtils.e(getName(), "Error processing class file: $fileName, error: ${e.message}")
                }
            }
        }
    }
}
